import numpy as np
import image_io
from image_utils import logging


def split_image_channels(image):
    '''
    该函数用于分割图像的颜色通道
    :param image: 输入图像，numpy数组，形状为(H,W,C)
    :return: 图像的颜色通道列表，每个元素形状为(H,W)
    '''
    try:
        logging.info("开始分割图像颜色通道")
        
        # 1. 参数验证
        if not isinstance(image, np.ndarray):
            raise TypeError("输入图像必须是numpy数组")
            
        # 2. 检测是否为多通道图像
        if len(image.shape) != 3 or image.shape[2] < 2:
            raise ValueError("输入必须是多通道图像(通道数≥2)")
        
        # 3. 分割通道并去除多余维度
        channels = np.split(image, image.shape[2], axis=2)
        channel_list = [np.squeeze(channel, axis=2) for channel in channels]
        
        logging.info(f"成功分割图像为 {len(channel_list)} 个通道")
        return channel_list
    except Exception as e:
        logging.error(f"分割图像通道失败: {e}")
        raise ValueError(f"在分割图像颜色通道时发生错误: {e}")


def merge_image_channels(file_paths):
    '''
    该函数用于合并多个单通道图像，创建一个多通道图像
    :param file_paths: 包含单通道图像文件路径的列表
    :return: 合并后的多通道图像，形状为(H,W,C)
    '''
    try:
        logging.info(f"开始合并 {len(file_paths)} 个图像通道")
        
        # 1. 参数验证
        if not file_paths or len(file_paths) < 2:
            raise ValueError("至少需要2个图像路径进行合并")
            
        # 2. 加载并验证图像
        base_shape = None
        channels = []
        
        for path in file_paths:
            img = image_io.read_image(path)
            
            # 验证是否为单通道
            if len(img.shape) > 2 and img.shape[2] > 1:
                raise ValueError(f"图像 {path} 必须是单通道图像")
                
            # 确保所有图像尺寸一致
            if base_shape is None:
                base_shape = img.shape[:2]
            elif img.shape[:2] != base_shape:
                raise ValueError("所有图像必须具有相同的尺寸")
                
            # 如果是灰度图但shape是(H,W)，调整为(H,W,1)
            if len(img.shape) == 2:
                img = np.expand_dims(img, axis=2)
                
            channels.append(img)
        
        # 3. 合并通道
        merged_image = np.concatenate(channels, axis=2)
        
        logging.info(f"成功合并为 {merged_image.shape[2]} 通道图像")
        return merged_image
    except Exception as e:
        logging.error(f"合并图像通道失败: {e}")
        raise ValueError(f"在合并图像通道时发生错误: {e}")